#!/usr/bin/python


"""

exon_structure.py

   holds intron exon structure


"""

from string import *
import sys, os
import StringIO

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))

# get module name
if __name__ == '__main__':
    my_module_name = os.path.split(sys.argv[0])[1]
    my_module_name = os.path.splitext(my_module_name)[0];
else:
    my_module_name = __name__

    
from general_util import die_error



from line_break import line_wrap

class exon:
    def __init__ (self, beg, end, coding_regions, coding_len, comment):
        self.beg            = beg          
        self.end            = end         
        self.coding_regions = coding_regions 
        self.coding_len     = coding_len
        self.comment        = comment        
        



from translate_by_codon import translate



        
class gene_exons (object):


    def __init__ (self, contig_or_chromosome, strand):
        """
        __init__

        """
        self.first_exon   = None
        self.exons        = list()
        self.introns      = list()
        self.contig       = contig_or_chromosome
        self.strand       = strand
        self.first_coding = None
        self.last_coding  = None



    #-------------------------------------------------------------------------------------
    #
    #   finish_init
    # 
    # 
    def finish_init (self):
        """
        finish initialising:
            calculate first coding and last coding
        """

        if not self.first_exon:
            return

        #
        #   set first and last coding exon
        # 
        # first exon
        if self.first_exon.coding_len:
            # this is the first coding exon
            self.first_coding   = -1
            # and maybe the last as well
            self.last_coding    = -1

            
        # other exons
        for index, exon in enumerate(self.exons):

            if exon.coding_len:
                if self.first_coding == None:
                    # this is the first coding exon
                    self.first_coding = index
                # and maybe the last coding exon
                self.last_coding = index


        # 
        # make sure no overlaps!!!
        # 
        prev_end = self.first_exon.end
        for index, exon in enumerate(self.exons):
            if exon.beg <= prev_end:
                raise "Exons #%d and #%d overlap" % (index + 1, index + 2)
            prev_end = exon.end
        
        
    #-------------------------------------------------------------------------------------
    #
    #   add_exon
    # 
    # 
    def add_exon (self, beg, end, orig_coding_regions, comment):
        """
            Add exon.
            If this is not the first exon, will add intron to list as well
        
        """
        from interval_list import intervals
        coding_regions = intervals(sorted(orig_coding_regions))
        coding_regions.combine_close_neighbours (1)

        # get total length of coding region
        coding_len = coding_regions.total_len()
            
        # assign to first exon or other exons
        if self.first_exon == None:
            self.first_exon = exon(beg, end, coding_regions.data, coding_len, comment)
        else:
            if len(self.exons):
                self.introns.append([self.exons[-1].end, beg])
            else:
                self.introns.append([self.first_exon.end, beg])
            self.exons.append(exon(beg, end, coding_regions.data, coding_len, comment))

            
    #8888888888888888888888888888888888888888888888888888888888888888888888888888888888888

    #        Per exon iterators
    #
    #           iter_len
    #           iter_cds_len
    #           iter_frames
    #           iter_cds
    #           iter_descriptions
    #           iter_translation
    #           iter_exons
    # 
    #           iter_intron_exons
    # 
    #        Per intron iterators
    # 
    #           iter_intron_len
    #           iter_splice_sites
    # 
    #       Fow whole gene
    # 
            
    #8888888888888888888888888888888888888888888888888888888888888888888888888888888888888


            
    #-------------------------------------------------------------------------------------
    #
    #   iter_len
    # 
    # 
    def iter_len (self):
        """
        returns lengths of exon
        """
        if self.first_exon:
            yield self.first_exon.end - self.first_exon.beg

        for exon in self.exons:
            yield exon.end - exon.beg
        
    #-------------------------------------------------------------------------------------
    #
    #   iter_cds_len
    # 
    def iter_cds_len (self):
        """
        returns lengths of CDS
        """
        if self.first_exon:
            yield self.first_exon.coding_len

        for exon in self.exons:
            yield exon.coding_len
        
    #-------------------------------------------------------------------------------------
    #
    #   iter_frames
    # 
    # 
    def iter_frames (self):
        """
        returns lengths of CDS
        """
        curr_frame = 0
        if self.first_exon:
            end_frame = (3 - (curr_frame + self.first_exon.coding_len) % 3) % 3
            if self.first_exon.coding_len:
                yield curr_frame, end_frame
            else:
                yield -1,-1
            curr_frame = end_frame

        for exon in self.exons:
            end_frame = (3 - (curr_frame + exon.coding_len) % 3) % 3
            if exon.coding_len:
                yield curr_frame, end_frame
            else:
                yield -1,-1
            curr_frame = end_frame


    #-------------------------------------------------------------------------------------
    #
    #   iter_cds
    # 
    # 
    def iter_cds(self, dna):
        def get_cds_of_specified_exon (dna, exon):
            return "".join(dna[exon.beg+beg:exon.beg+end].upper()
                                    for beg, end in exon.coding_regions)
        if not self.first_exon:
            return

        yield(get_cds_of_specified_exon(dna, self.first_exon))

        for exon in self.exons:
            yield(get_cds_of_specified_exon(dna, exon))

    #-------------------------------------------------------------------------------------
    #
    #   iter_intron_len
    # 
    # 
    def iter_intron_len (self):
        """
        returns lengths of introns
        """
        for b, e in self.introns:
            yield e - b

    #-------------------------------------------------------------------------------------
    #
    #   iter_splice_sites
    # 
    # 
    def iter_splice_sites (self, dna):
        """
        returns intron splice sites
        """
        for b, e in self.introns:
            if e - b < 4:
                yield ("","")
            else:
                yield (dna[b:b+2].lower(), dna[e-2:e].lower())

    #-------------------------------------------------------------------------------------
    #
    #   iter_translation
    # 
    # 
    def iter_translation (self, dna, frame_straddling):
        """
        returns peptide per exon
        """

        peptide = self.get_gene_translation(dna)

        cummulative_cds_len = 0
        for cds_len, (frame_beg, frame_end) in izip(self.iter_cds_len(), self.iter_frames()):

            # ignore non coding
            if not cds_len:
                yield ""
                continue

            # translation
            exon_translation = peptide[(cummulative_cds_len+2) // 3 : 
                                      (cummulative_cds_len + cds_len +2) // 3]

            # add frame around straddling residues
            if frame_end != 0 and frame_straddling:
                exon_translation = exon_translation[:-1] + "[" + exon_translation[-1:] + "]"

            yield exon_translation

            cummulative_cds_len += cds_len
            
    #-------------------------------------------------------------------------------------
    #
    #   iter_exons
    # 
    # 
    def iter_exons(self, dna):
        """
            retrieves exons
            
        """
        if not self.first_exon:
            return
        
            
        # 
        # First exon
        # 
        yield  self.exon_get_sequence(dna, self.first_exon)
        

        # 
        # Other exons
        # 
        for exon in self.exons:
            yield self.exon_get_sequence(dna, exon)

        


    #-------------------------------------------------------------------------------------
    #
    #   iter_descriptions
    # 
    # 
    def iter_descriptions (self):
        """
            retrieves descriptions of exons
            
        """
        if not self.first_exon:
            return
        
        # 
        # First exon
        # 
        running_totals = dict(curr_frame=0, cds_total_len= 0)
        description = self.exon_get_description (self.first_exon, 
                                                   self.first_coding == -1,
                                                   self.last_coding == -1, 
                                                   1, running_totals)
        yield description
        
        # 
        # Other exons
        # 
        for exon_index, exon in enumerate(self.exons):
            # exon index ++ for 1-based and another ++ for first_exon
            description = self.exon_get_description(exon,
                                                    exon_index <= self.first_coding,
                                                    exon_index  >= self.last_coding , 
                                                    exon_index + 2, running_totals)
            yield description
        
            
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#               Per gene                  
                  
    #-------------------------------------------------------------------------------------
    #
    #   get_gene_cds
    # 
    # 
    def get_gene_cds(self, dna):
        return "".join(self.iter_cds(dna))

    #-------------------------------------------------------------------------------------
    #
    #   get_gene_translation
    # 
    # 
    def get_gene_translation(self, dna):
        return translate(self.get_gene_cds(dna))

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#               Per exon helpers
                  


    #-------------------------------------------------------------------------------------
    #
    #   get_exon_description
    # 
    # 
    def exon_get_description (self, exon, utr5, utr3, exon_index, running_totals):
        """
        returns description of exon, including:
            beg, end and length
            and if CDS
                beg and end frames
                beg and end coding
                coding length
                total coding length
        """
        exon_length = exon.end - exon.beg
        descriptions = ["exon #%d: %d-%d" % (exon_index, exon.beg, exon.end)]
        descriptions.append("len = %d" % exon_length)

        # 5' UTR
        if utr5:
            if not exon.coding_regions:
                descriptions.append("5'UTR_len = %d" % exon_length)
            #elif exon.coding_regions[0][0] != 0:
            else:
                descriptions.append("5'UTR_len = %d" % exon.coding_regions[0][0])

        # CDS
        if exon.coding_len:

            # CDS len and total len
            descriptions.append("CDS_len = %d" % exon.coding_len)
            running_totals['cds_total_len'] += exon.coding_len
            descriptions.append("CDS_total_len = %d" %  + running_totals['cds_total_len'])

            # Frame
            end_frame = (3 - (running_totals['curr_frame'] + exon.coding_len) % 3) % 3
            descriptions.append("Frame = (%d-%d)" % (running_totals['curr_frame'], end_frame))
            running_totals['curr_frame'] = end_frame
        
        # 3' UTR
        if utr3:
            if not exon.coding_regions:
                descriptions.append("3'UTR_len = %d" % exon_length)
            else:
                cds_end = exon.coding_regions[-1][1]
                descriptions.append("3'UTR_len = %d" % (exon.end - exon.beg - cds_end))

        return "; ".join (descriptions) 
        


        

    #-------------------------------------------------------------------------------------
    #
    #   get_sequence_of_specified_exon
    # 
    # 
    def exon_get_sequence(self, dna, exon):
        prev_end = exon.beg
        seq_segments = list()
        for beg, end in exon.coding_regions:
            seq_segments.append(dna[prev_end:exon.beg+beg].lower())
            seq_segments.append(dna[exon.beg+beg:exon.beg+end].upper())
            prev_end = exon.beg+end
        
        seq_segments.append(dna[prev_end:exon.end].lower())
            
        return "".join(seq_segments)  + "\n%d-%d%s" % (exon.beg, exon.end, str(exon.coding_regions))
        
    #-------------------------------------------------------------------------------------
    #
    #   get_intron_exon_sequence
    # 
    # 
    def iter_intron_exons (self, dna, prepostgenic_limit = 0):
        """
            retrieves intron and exons sequence to specified file

            yields description as well because 
            iter_descriptions only gives exon and not intron descriptions

            prepostgenic_limit is how much of the up and down stream gene to collect
            
        """
        if not self.first_exon:
            return
        
        # 
        # Intergenic before beg
        # 
        if prepostgenic_limit != 0 and self.first_exon.beg != 0:
            pre_end = self.first_exon.beg
            pre_beg = max(0,  pre_end - prepostgenic_limit)
            description = "intergenic %d-%d; len = %d" % \
                            (pre_beg, pre_end, pre_end - pre_beg)
            yield (description, dna[pre_beg:self.first_exon.beg].lower())


        running_totals = dict(curr_frame=0, cds_total_len= 0)

            
        # 
        # First exon
        # 
        description = self.exon_get_description (self.first_exon, 
                                                   self.first_coding == -1,
                                                   self.last_coding == -1, 
                                                   1, running_totals)
        exon_sequence = self.exon_get_sequence(dna, self.first_exon)
        yield(description, exon_sequence)
        

        # 
        # Other exons
        # 
        prev_end = self.first_exon.end
        for exon_index, exon in enumerate(self.exons):
            #
            # intron
            # 
            description = "intron #%d %d-%d; len = %d" % \
                                    (exon_index + 1, prev_end, exon.beg, exon.beg - prev_end)
            yield(description, dna[prev_end:exon.beg].lower())
            prev_end = exon.end

            
            #
            # exon
            # 
            # exon index ++ for 1-based and another ++ for first_exon
            description = self.exon_get_description(exon,
                                                    exon_index <= self.first_coding,
                                                    exon_index  >= self.last_coding , 
                                                    exon_index + 2, running_totals)
            exon_sequence = self.exon_get_sequence(dna, exon)
            yield(description, exon_sequence)

        
        # 
        # Intergenic after end
        # 
        if prepostgenic_limit != 0 and prev_end < len(dna):
            post_end = min(prev_end + prepostgenic_limit, len(dna))
            description = "postgenic %d-%d; len = %d" % \
                            (prev_end, post_end, post_end - prev_end)
            yield (description, dna[prev_end:post_end].lower())
            



    #-------------------------------------------------------------------------------------
    #
    #   __str__
    # 
    # 
    def __str__ (self):
        """
            for debugging
            
        """
        if not self.first_exon:
            return "Empty gene structure"
        return (
                        "#" + self.get_summary() +
                            "\n#1 %d-%d (cds=%d) %s" % 
                                (self.first_exon.beg,
                                 self.first_exon.end,
                                 self.first_exon.coding_len,
                                 self.first_exon.comment) + "\n" + 
        
                            ("\n".join("#%d %d-%d (cds=%d) %s" % 
                                (i +2, e.beg, e.end, e.coding_len, e.comment) 
                                for i, e in enumerate(self.exons))))

                    
    #-------------------------------------------------------------------------------------
    #
    #   get_summary
    # 
    # 
    def get_summary (self):
        """
            return summary
            
        """

        if not self.first_exon:
            return "Empty gene structure"

        #total length
        total_length = self.first_exon.end - self.first_exon.beg
        total_length += sum(e.end - e.beg for e in self.exons)

        # total coding length
        total_coding_length = self.first_exon.coding_len
        total_coding_length += sum(e.coding_len for e in self.exons)

        exon_count = len(self.exons) + 1
        
        return ("contig=%s;"\
                    "strand=%s;"\
                    "%d exons;"\
                    "total_coding_length=%d;"\
                    "total_length=%d;"\
                    "CDS_beg=#%d;"\
                    "CDS_end=#%d;"          % (
                                            self.contig, 
                                            self.strand, 
                                            exon_count,
                                            total_coding_length,
                                            total_length,
                                            self.first_coding,
                                            self.last_coding))

        

    #-------------------------------------------------------------------------------------
    #
    #   tabulated_summary
    # 
    # 
    def tabulated_summary_columns (self):
        return ("#intron","5'-splice","3'-splice","frame","frame","len","cds")
    
    def tabulated_summary (self, dna):

        if not self.first_exon:
            return

        # write out intron exon structure
        exon_lengths, cds_lengths, frames  = (self.iter_len(), 
                                      self.iter_cds_len(),
                                      self.iter_frames())
    
        # first exon
        frame_b, frame_e = frames.next()
        if frame_b == -1:
            frame_b = '-'
            frame_e = '-'
        else:
            frame_b, frame_e = str(frame_b), str(frame_e)
        cds_len = cds_lengths.next()
        if cds_len == 0:
            cds_len = '-'
        else:
            cds_len = str(cds_len)

        # for first exon
        yield ["","","", frame_b, frame_e, str(exon_lengths.next()), cds_len]
    
        #introns and other exons
        for intron_len, (splice5, splice3), exon_len, cds_len, (frame_b, frame_e) in \
                                    izip( self.iter_intron_len(),
                                          self.iter_splice_sites (dna),
                                          exon_lengths, cds_lengths, frames):
            if frame_b == -1:
                frame_b = '-'
                frame_e = '-'
            else:
                frame_b, frame_e = str(frame_b), str(frame_e)
            if cds_len == 0:
                cds_len = '-'
            else:
                cds_len = str(cds_len)
            yield [str(intron_len), splice5, splice3, frame_b, frame_e, 
                   str(exon_len), cds_len]
        
    #-------------------------------------------------------------------------------------
    #
    #   load_from_fasta
    # 
    #       returns both gene structure and DNA in one string
    # 
    # 
    def load_from_fasta(seq_lines):
        gene = gene_exons("from file", "+")

        from fasta_sequences_iterator import fasta_sequences_iterator

        import string
        allchars = string.maketrans('','')
        allbutacgt = allchars.translate(allchars, 'ACGTUMRSVWYHKDBNacgtumrsvwyhkdbn-~')

        import re
        coding_re = re.compile("([^a-z]+)")
        
        curr_pos = 0
        cummulative_dna = list()
        for acc, seq in fasta_sequences_iterator(seq_lines):

            # remove not acgt etc.
            seq = seq.translate(allchars, allbutacgt)
            cummulative_dna.append(seq)
            if not len(seq):
                continue

            
            if "exon" in acc:
                coding_regions = list()
                for m in coding_re.finditer(seq):
                    coding_regions.append(m.span())
                exon_len = len(seq)
                gene.add_exon(curr_pos, 
                                curr_pos + exon_len,
                                coding_regions, acc)
                curr_pos += exon_len
            else:
                curr_pos += len(seq)
        gene.finish_init()

        return ("".join(cummulative_dna), gene)
    load_from_fasta = staticmethod(load_from_fasta)

                        


from itertools import *
def unit_test_helper(gene,dna):
    #print gene

    sys.stdout.write ("\n\n-------Genomic exons-------\n")
    sys.stdout.write("#" + gene.get_summary() + "\n")
    sys.stdout.write("#" + str(gene) + "\n")
    for description, seq in gene.iter_intron_exons (dna, 50000):
        sys.stdout.write( ">" + description + "\n" + seq + "\n")
    sys.stdout.write ("\n-------Exon structure-------\n")
    line_format = "%-9s " * 7  + "\n"
    sys.stdout.write(line_format % gene.tabulated_summary_columns())
    for cols in gene.tabulated_summary(dna):
        sys.stdout.write(line_format % tuple(cols))

    for description, peptide in izip(gene.iter_descriptions(), gene.iter_translation (dna, True)):
        if not len(peptide):
            continue
        sys.stdout.write(">" + description + "\n" + peptide + "\n")
    print "\n"
    return


def unit_test1():
    #pppppppppppppppppppppp
    #uuu
    #gtiiiag
    #uuuuuuuuuuAGTAGTAGTAGTAGTAGTAGTAGTA
    #gtiiiiiiiiiiiiiiiiag
    #GTAGTAGnnnTAGTAGT
    #gtiiiiiiiiiiiag
    #AGTAGTAGTAGTuuuuuuuuuu
    #gtiiiag
    #uuu
    #ppppppppp
    # 0  :22  
    # 22 :57    10:35
    # 57 :77    
    # 77 :94    0:7 10:17
    # 94 :109   
    # 109:131  0:12
    # 131:140
    dna = "ppppppppppppppppppppppuuugtiiiaguuuuuuuuuuAGTAGTAGTAGTAGTAGTAGTAGTAgtiiiiiiiiiiiiiiiiagGTAGTAGnnnTAGTAGTgtiiiiiiiiiiiagAGTAGTAGTAGTuuuuuuuuuugtiiiaguuuppppppppp"
    gene = gene_exons("chromosome1", '+')
    #gene.add_exon("chromosome1", '+', 0  , 22 , [], "pregenic")
    gene.add_exon(22 , 25 , [], "5'utr")
    gene.add_exon(32 , 67 , [(10, 35)], "first exon")
    gene.add_exon(87 , 104 , [(0,7), (10,17)], "second exon with break")
    gene.add_exon(119, 141, [(0,12)], "third exon")
    gene.add_exon(148, 151, [], "3'utr")
    gene.finish_init()
    unit_test_helper(gene, dna)

def unit_test2():
    fasta = \
""">intergenic
nnnnnnnnnnnnnnnnnnnnnn
>exon #1 5utr
uuu
>intron1
gtvvvag
>exon #2 5utr
uuuuuuuuuuAGTAGTAGTAGTAGTAGTAGTAGTA
>intron2
gtvvvvvvvvvvvvvvvvag
>exon #3
GTAGTAGnnnTAGTAGT
>intron3
gtvvvvvvvvvvvag
>exon #4
AGTAGTAGTAGTuuuuuuuuuu
>intron4
gtvvvag
>exon #5
uuu
>intergenic
nnnnnnnnn"""
    dna = "ppppppppppppppppppppppuuugtiiiaguuuuuuuuuuAGTAGTAGTAGTAGTAGTAGTAGTAgtiiiiiiiiiiiiiiiiagGTAGTAGnnnTAGTAGTgtiiiiiiiiiiiagAGTAGTAGTAGTuuuuuuuuuugtiiiaguuuppppppppp"
    dna, gene = gene_exons.load_from_fasta(fasta.split("\n"))
    unit_test_helper(gene, dna)


if __name__ == '__main__':
    unit_test1()    
    unit_test2()    



